from contextlib import contextmanager
from colorsys import hsv_to_rgb
from bisect import bisect_left

from utils.interval import Interval
from utils.plotter import get_plot_target

from geom2d.segment import Segment

from cranesched.routing.path import Path
from cranesched.routing.bounds import compute_bound, UnavoidableCollision


class Bounds(object):
    """Associates a Movable1D with the lower and upper bounds on its movement."""
    def __init__(self, movable):
        self.movable = movable
        self.lb = None
        self.ub = None

    def __iter__(self):
        """This allows 'lb, ub = bounds'."""
        return iter((self.lb, self.ub))

    def compute_lb(self, entry_below, force=False):
        if force or self.lb is None:
            distance = (self.movable.size + entry_below.movable.size) * 0.5
            self.lb = compute_bound(self.movable.path, entry_below.lb,
                                    direction=-1, speed=self.movable.speed,
                                    distance=distance)

    def compute_ub(self, entry_above, force=False):
        if force or self.ub is None:
            distance = (self.movable.size + entry_above.movable.size) * 0.5
            self.ub = compute_bound(self.movable.path, entry_above.ub,
                                    direction=+1, speed=self.movable.speed,
                                    distance=distance)

    def is_out_of_bounds(self, time, position):
        """Return a boolean indicating whether 'position' is not between the movable's LB and UB at
        the given 'time'."""
        return not (self.lb.position_at(time) <= position <= self.ub.position_at(time))

    def is_reachable(self, position):
        """True if 'position' can be reached at some point in the future."""
        return self.lb.curr_position <= position <= self.ub.curr_position


class DummyMovable(object):
    """Just an object that should behave as a Movable1D object to the eyes of the router, but is
    otherwise featureless."""

WALL = DummyMovable()
WALL.size = 0.0

MAX_T = 1e10


class Router(object):
    """This object can be used to plan valid paths (described by line segments) for satisfying
    movement requests while avoiding collisions between Movable1D objects along a specific
    movement axis.
    In order to find valid (i.e. collision-less) routes, the router keeps an association between
    each Movable1D object and its movement bounds. This object is responsible for maintaining
    updated lower and upper bounds for the movement of each Movable1D object, taking their
    previously planned paths into account."""
    UnavoidableCollision = UnavoidableCollision
    MAX_T = MAX_T

    def __init__(self, lo=0.0, hi=100.0, time=0.0):
        self.limits = Bounds(WALL)
        self.limits.lb = Path(time, lo)
        self.limits.lb.append(MAX_T, lo)
        self.limits.ub = Path(time, hi)
        self.limits.ub.append(MAX_T, hi)
        self.movable_to_bounds = {}  # {Movable1D: Bounds}
        self.movables = []           # [Movable1D] (sorted by position at initial time)
        self.stack = []              # stack of saved bounds mappings
        self.initialized = False     # bool (router allows certain operations only when
                                     # uninitialized, or others only when initialized)

    def __getitem__(self, i):
        return self.movables[i]

    def __iter__(self):
        return iter(self.movables)

    def get_initial_time(self):
        return self.limits.lb.times[0]

    def set_initial_time(self, time):
        if self.initialized:
            raise Exception("set_initial_time() can only be called on uninitialized router")
        self.limits.lb.shift_time(time - self.limits.lb.times[0])
        self.limits.ub.shift_time(time - self.limits.ub.times[0])

    def set_limits(self, lo=None, hi=None):
        """Sets the "walls" between which the Movable1D objects can move."""
        if self.initialized:
            raise Exception("set_limit() can only be called on uninitialized router")
        if lo is not None:
            self.limits.lb.shift_position(lo - self.limits.lb.curr_position)
        if hi is not None:
            self.limits.ub.shift_position(hi - self.limits.ub.curr_position)

    def stretch_limits(self, factor):
        """Stretch the limits set by router's "walls" by 'factor' around its center."""
        lo = self.limits.lb.curr_position
        hi = self.limits.ub.curr_position
        center = (lo + hi) * 0.5
        amplitude = (hi - lo) * factor
        self.set_limits(center - amplitude * 0.5, center + amplitude * 0.5)

    def add(self, movable):
        """Adds a Movable1D object to the router. The object is inserted into the movables sequence,
        ordered by its position at the router's initial time."""
        if self.initialized:
            raise Exception("add() can only be called on uninitialized router")
        initial_time = self.get_initial_time()
        initial_positions = [m.path.position_at(initial_time) for m in self.movables]
        index = bisect_left(initial_positions, movable.path.position_at(initial_time))
        self.movables.insert(index, movable)

    def init(self):
        """Initialize the router. After this point, the router disables addition of new movables,
        setting limits (i.e. walls), or initial time. The router permits only query operations and
        update() calls when initialized."""
        if self.initialized:
            raise Exception("router is already initialized")
        # self._validate_initial_state()
        self.initialized = True
        self.clear()

    # def _validate_initial_state(self):
    #     """This method checks if the paths of the movables passed to the router are valid, that is
    #     the paths are do not intersect."""
    #     init_time = self.get_initial_time()
    #     for i in xrange(len(self.movables) - 1):
    #         curr = self.movables[i]
    #         next = self.movables[i+1]
    #         dist = (curr.size + next.size) * 0.5
    #         temp = deepcopy(next.path)
    #         temp.shift_position(-dist)
    #         curr_init_pos = curr.path.position_at(init_time)
    #         temp_init_pos = temp.position_at(init_time)
    #         if curr_init_pos > temp_init_pos:
    #             raise Exception("neighbor movables are too close or reversed")
    #         if curr.path.crosses(next.path):
    #             raise Exception("planned paths of neighbor movables contain a collision")

    def push(self):
        """Saves the current state of the bounds mapping into the stack."""
        if not self.initialized:
            raise Exception("router is not initialized")
        self.stack.append(self.movable_to_bounds)
        self.clear()

    def pop(self):
        """Reverts the bounds mapping to the one that last saved."""
        if not self.initialized:
            raise Exception("router is not initialized")
        self.movable_to_bounds = self.stack.pop()

    def clear(self):
        """Creates new (empty) bounds entries for all movables in the router."""
        if not self.initialized:
            raise Exception("router is not initialized")
        self.movable_to_bounds = {movable: Bounds(movable) for movable in self.movables}

    def update(self, force=False):
        """Computes movement lower and upper bounds for all Movable1D objects."""
        if not self.initialized:
            raise Exception("router is not initialized")
        if force:
            self.clear()
        bounds = self.movable_to_bounds
        # compute lower bounds from "bottom" to "top"
        below = self.limits
        for movable in self.movables:
            current = bounds[movable]
            current.compute_lb(below)
            below = current
        # compute upper bounds from "top" to "bottom"
        above = self.limits
        for movable in reversed(self.movables):
            current = bounds[movable]
            current.compute_ub(above)
            above = current

    def is_reachable(self, movable, position):
        """Return true if 'movable' can at some arbitrary point in the future reach 'position'."""
        bounds = self.movable_to_bounds[movable]
        return bounds.is_reachable(position)

    def earliest_fit(self, movable, target, duration=0.0, window=None, verbose=False):
        """Finds the earliest possible fit (considering the object's movement bounds) of the given
        request to move 'movable' into position 'target', and stay in that position for 'duration'
        seconds within interval 'window' (default window=[-MAX_T, +MAX_T], and MAX_T=1e10). Returns
        the start time of the interval which satisfies the request, i.e. the time when the object
        reaches position 'target'."""
        if not self.initialized:
            raise Exception("router is not initialized")
        log = do_print if verbose else no_print
        if window is None:
            window = Interval(self.get_initial_time(), MAX_T)
        earliest_start = max(window.start, movable.earliest_arrival_at(target))
        if earliest_start + duration > window.end:
            # even if we moved immediately into the target position, we wouldn't be
            # able to complete the action before the end of the satisfiability window
            log("fit unavailable")
            return None

        bounds = self.movable_to_bounds[movable]
        oob = bounds.is_out_of_bounds(earliest_start, target)
        t0 = t1 = earliest_start
        log("out of bounds=%s, interval=[%s, %s]", oob, t0, t1)
        if earliest_start == window.end:
            # surely 'duration' is zero, otherwise the condition
            # on the previous if statement would be true
            assert duration == 0.0
            if oob:
                log("fit unavailable")
                return None
            else:
                log("earliest fit=%s", earliest_start)
                return earliest_start

        eligible = Segment((earliest_start, target), (window.end, target))
        intersections = Path.multipath_segment_intersections((bounds.lb, bounds.ub), eligible)
        for (t2, _), _, _ in intersections:
            log("intersection at t=%s", t2)
            if t1 == t2:
                log("skipped")
                continue
            oob, t0, t1 = update_out_of_bounds(oob, t0, t1, t2, target, bounds)
            log("out of bounds=%s, interval=[%s, %s]", oob, t0, t1)
            if not oob and t1 - t0 >= duration:
                log("earliest fit=%s", t0)
                return t0
        if t1 < window.end:
            oob, t0, t1 = update_out_of_bounds(oob, t0, t1, window.end, target, bounds)
        log("out of bounds=%s, interval=[%s, +INF)", oob, t0)
        if not oob and t1 - t0 >= duration:
            log("earliest fit=%s", t0)
            return t0
        log("fit unavailable")
        return None

    def compute_trajectory(self, movable, target, time):
        """Computes a trajectory for a 'movable' to reach a 'target' position by a specified 'time'.
        The given target position must obviously lie between the lower and upper bounds for the
        movable's position. Returns a sequence of points that should be appended to the Movable's
        path (using extend() for example, or iterating over the list and using append())."""
        if not self.initialized:
            raise Exception("router is not initialized")
        lb, ub = self.movable_to_bounds[movable]
        path = movable.path
        points = [(time, target)]
        count = 0
        # Loop until we join with the movable's current time
        assert time >= path.curr_time
        while time > path.curr_time:
            assert len(points) == count + 1
            count += 1
            time, target = points[-1]

            # 1) See if we're at one of the bounds, and follow it if necessary.
            if time <= lb.curr_time:
                lb_segment = lb.segments_at(time)[0]
                if lb_segment.y_at(time) == target and lb_segment.m <= 0.0:
                    points.append(lb_segment.p0)
                    continue
            if time <= ub.curr_time:
                ub_segment = ub.segments_at(time)[0]
                if ub_segment.y_at(time) == target and ub_segment.m >= 0.0:
                    points.append(ub_segment.p0)
                    continue
            # 2) Otherwise, we're between the two bounds, so we need
            #    to approach one of them with a horizontal segment.
            horizontal_segment = Segment((time, target), (path.curr_time, target))
            iterator = Path.multipath_segment_intersections((lb, ub),
                                                            horizontal_segment,
                                                            reverse=True)
            for (t, p), _, _ in iterator:
                if t < time:
                    points.append((t, p))
                    break

        assert path.curr_time == time
        assert path.curr_position == target
        points.pop()
        points.reverse()
        return points

    def show(self, target=None, title=None):
        axes = self.show_planned_routes(target, title)
        self.show_movement_bounds(axes, None)
        self.show_walls(axes, None)
        return axes

    def show_planned_routes(self, target=None, title=None):
        with _plot_target(target, title) as axes:
            for i, movable in enumerate(self.movables):
                # compute a color for the movable, evenly distributed along the hue spectrum
                color = hsv_to_rgb(float(i) / len(self.movables), 1, 1)
                # plot the planned path first
                movable.path.line_plot(target=axes, color=color, linewidth=2)
                # fill area representing the movable's width
                offset = movable.size * 0.5
                below = [pos - offset for pos in movable.path.positions]
                above = [pos + offset for pos in movable.path.positions]
                axes.fill_between(movable.path.times, below, above, color=color, alpha=0.5)
        return axes

    def show_movement_bounds(self, target=None, title=None):
        with _plot_target(target, title) as axes:
            for i, movable in enumerate(self.movables):
                # compute a color for the movable, evenly distributed along the hue spectrum
                color = hsv_to_rgb(float(i) / len(self.movables), 1, 1)
                # fill area between lower and upper bound
                try:
                    lb, ub = self.movable_to_bounds[movable]
                except KeyError:
                    continue
                if lb is None or ub is None:
                    continue
                times = sorted(t for t in set(lb.times+ub.times) if t >= movable.path.curr_time)
                if times[-1] == MAX_T:
                    times.pop()
                lb_pos = map(lb.position_at, times)
                ub_pos = map(ub.position_at, times)
                axes.fill_between(times, lb_pos, ub_pos, color=color, alpha=0.25)
                axes.plot(times, lb_pos, linestyle="--", color=color)
                axes.plot(times, ub_pos, linestyle=":", color=color)
        return axes

    def show_walls(self, target=None, title=None):
        with _plot_target(target, title) as axes:
            max_t = 0.0
            for movable in self.movables:
                max_t = max(max_t, movable.path.curr_time)
                try:
                    bounds = self.movable_to_bounds[movable]
                except KeyError:
                    continue
                if bounds.lb and bounds.ub:
                    max_t = max(max_t, bounds.lb.times[-2], bounds.ub.times[-2])
            for wall in self.limits:
                axes.plot([0, max_t], [wall.curr_position, wall.curr_position],
                          color="black", linestyle="dashed")


@contextmanager
def _plot_target(target, title=None):
    plt = get_plot_target(target)
    axes = plt.get_axes()
    yield axes
    if title is not None:
        axes.set_title(str(title))
    axes.grid(True)
    plt.update()


def update_out_of_bounds(oob, t0, t1, t2, pos, bounds):
    """Auxiliary function used by Router.earliest_fit()."""
    if oob != bounds.is_out_of_bounds((t1 + t2) * 0.5, pos):
        return not oob, t1, t2
    return oob, t0, t2


def do_print(message, *args):
    print message % args


def no_print(message, *args):
    pass
